package com.tony.ai.ai

import android.content.Context
import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.graphics.Color
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.segmentation.Segmentation
import com.google.mlkit.vision.segmentation.SegmentationMask
import com.google.mlkit.vision.segmentation.Segmenter
import com.google.mlkit.vision.segmentation.selfie.SelfieSegmenterOptions
import org.opencv.android.Utils
import org.opencv.core.Mat
import org.opencv.core.Size
import org.opencv.imgproc.Imgproc

/**
 * 图像分割工具
 */
class SegmentationReplacer {

    enum class ReplacerType {
        COLOR,
        IMAGE,
        BACKGROUND,
        TRANSPARENT,//背景透明
    }

    private val segmenter: Segmenter

    init {
        // 初始化 ML Kit 分割器
        val options = SelfieSegmenterOptions.Builder()
            .setDetectorMode(SelfieSegmenterOptions.STREAM_MODE) // 适用于实时处理
            .build()
        segmenter = Segmentation.getClient(options)
    }

    /**
     * 替换背景
     * @param originalBitmap 原始图像
     * @param backgroundBitmap 新背景图像
     * @param onSuccess 成功回调，返回替换背景后的图像
     * @param onFailure 失败回调
     */
    private fun replaceBackground(
        type: ReplacerType,
        originalBitmap: Bitmap,
        backgroundBitmap: Bitmap,
        color: Int,
        onSuccess: (Bitmap) -> Unit,
        onFailure: (Exception) -> Unit
    ) {
        val newOriginalBitmap = scaleBitmapHalf(originalBitmap)
        val inputImage = InputImage.fromBitmap(newOriginalBitmap, 0)
        segmenter.process(inputImage)
            .addOnSuccessListener { segmentationMask ->
                // 获取分割掩码并替换背景
                val resultBitmap =
                    applyMaskOptimize(
                        type,
                        newOriginalBitmap,
                        backgroundBitmap,
                        segmentationMask,
                        color
                    )
                onSuccess(resultBitmap)
            }
            .addOnFailureListener { e ->
                onFailure(e)
            }
    }


    /**
     * 应用分割掩码并替换背景
     */
    private fun applyMask(
        type: ReplacerType,
        originalBitmap: Bitmap,
        backgroundBitmap: Bitmap,
        segmentationMask: SegmentationMask,
        color: Int
    ): Bitmap {
        val mask = segmentationMask.buffer
        val maskWidth = segmentationMask.width
        val maskHeight = segmentationMask.height
        // 绘制新背景
        val scaledBackground =
            Bitmap.createScaledBitmap(backgroundBitmap, maskWidth, maskHeight, true)
        val maskBitmap = Bitmap.createBitmap(maskWidth, maskHeight, Bitmap.Config.ARGB_8888)
        for (y in 0 until maskHeight) {
            for (x in 0 until maskWidth) {
                val maskValue = mask.float
                val alpha = (maskValue * 255).toInt()
                if (alpha > 128) {
                    val pixel = originalBitmap.getPixel(x, y)
                    // 提取 ARGB 分量
                    val alpha = Color.alpha(pixel)
                    val red = Color.red(pixel)
                    val green = Color.green(pixel)
                    val blue = Color.blue(pixel)
                    maskBitmap.setPixel(x, y, Color.argb(alpha, red, green, blue))
                } else {
                    if (type == ReplacerType.COLOR) {
                        maskBitmap.setPixel(x, y, color)
                    } else {
                        val pixel = scaledBackground.getPixel(x, y)
                        // 提取 ARGB 分量
                        val alpha = Color.alpha(pixel)
                        val red = Color.red(pixel)
                        val green = Color.green(pixel)
                        val blue = Color.blue(pixel)
                        maskBitmap.setPixel(x, y, Color.argb(alpha, red, green, blue))
                    }
                }
            }
        }
        return maskBitmap
    }

    private fun applyMaskOptimize(
        type: ReplacerType,
        originalBitmap: Bitmap,
        backgroundBitmap: Bitmap,
        segmentationMask: SegmentationMask,
        color: Int
    ): Bitmap {
        val mask = segmentationMask.buffer
        val maskWidth = segmentationMask.width
        val maskHeight = segmentationMask.height

        // 缩放背景图像以匹配掩码分辨率
        val scaledBackground =
            Bitmap.createScaledBitmap(backgroundBitmap, maskWidth, maskHeight, true)

        // 创建结果 Bitmap
        val maskBitmap = Bitmap.createBitmap(maskWidth, maskHeight, Bitmap.Config.ARGB_8888)

        // 获取像素数组
        val originalPixels = IntArray(maskWidth * maskHeight)
        originalBitmap.getPixels(originalPixels, 0, maskWidth, 0, 0, maskWidth, maskHeight)

        val backgroundPixels = IntArray(maskWidth * maskHeight)
        scaledBackground.getPixels(backgroundPixels, 0, maskWidth, 0, 0, maskWidth, maskHeight)

        val resultPixels = IntArray(maskWidth * maskHeight)

        // 遍历每个像素
        for (y in 0 until maskHeight) {
            for (x in 0 until maskWidth) {
                val index = y * maskWidth + x

                // 获取掩码值
                val maskValue = mask.float
                val alpha = (maskValue * 255).toInt()

                // 判断是否属于前景
                if (alpha > 128) {
                    // 使用原始图像的像素
                    resultPixels[index] = originalPixels[index]
                } else {
                    // 根据类型设置背景
                    resultPixels[index] = when (type) {
                        ReplacerType.COLOR -> color
                        ReplacerType.BACKGROUND -> backgroundPixels[index]
                        ReplacerType.TRANSPARENT -> Color.TRANSPARENT
                        else -> color
                    }
                }
            }
        }

        // 将结果像素数组设置到 Bitmap 中
        maskBitmap.setPixels(resultPixels, 0, maskWidth, 0, 0, maskWidth, maskHeight)

        return maskBitmap
    }

    /**
     * 释放资源
     */
    fun release() {
        segmenter.close()
    }

    private fun scaleBitmapHalf(originalBitmap: Bitmap): Bitmap {
        // 计算新的宽度和高度，将原始宽度和高度分别除以 2
        val newWidth = originalBitmap.width / 2
        val newHeight = originalBitmap.height / 2
        // 创建缩放后的 Bitmap
        return Bitmap.createScaledBitmap(originalBitmap, newWidth, newHeight, true)
    }


    fun replaceBackgroundColor(
        originalImage: Bitmap, backgroundBitmap: Bitmap, isMatting: Boolean, color: Int,
        onResultSuccess: (resultBitmap: Bitmap) -> Unit
    ) {
        val type = if (isMatting) {
            ReplacerType.TRANSPARENT
        } else {
            ReplacerType.COLOR
        }
        replaceBackground(type,
            originalImage!!,
            backgroundBitmap!!,
            color,
            onSuccess = { resultBitmap ->
                onResultSuccess(resultBitmap)
            },
            onFailure = { e ->
                e.printStackTrace()
            }
        )
    }

    fun replaceBackgroundImage(
        mContext: Context,
        originalImage: Bitmap?,
        color: Int,
        isChecked: Boolean,
        isMatting: Boolean,
        resId: Int,
        onResultSuccess: (resultBitmap: Bitmap) -> Unit
    ) {
        var backgroundBitmap = BitmapFactory.decodeResource(mContext.resources, resId)
        val mBitmap = blurBitmap(backgroundBitmap, isChecked)
        if (mBitmap != null) {
            backgroundBitmap = mBitmap
        }
        val type = if (isMatting) {
            ReplacerType.TRANSPARENT
        } else {
            ReplacerType.BACKGROUND
        }
        replaceBackground(type,
            originalImage!!,
            backgroundBitmap!!,
            color,
            onSuccess = { resultBitmap ->
                onResultSuccess(resultBitmap)
            },
            onFailure = { e ->
                e.printStackTrace()
            }
        )
    }

    private fun blurBitmap(backgroundBitmap: Bitmap, isChecked: Boolean): Bitmap? {
        if (isChecked) {
            val mat = Mat()
            Utils.bitmapToMat(backgroundBitmap, mat)
            // 进行高度高斯模糊处理
            // 较大的卷积核尺寸和标准差可以实现高度模糊
            val kernelSize = 51 // 卷积核大小，必须为奇数
            val sigmaX = 20.0 // X 方向的高斯核标准差
            val sigmaY = 20.0 // Y 方向的高斯核标准差
            Imgproc.GaussianBlur(
                mat,
                mat,
                Size(kernelSize.toDouble(), kernelSize.toDouble()),
                sigmaX,
                sigmaY
            )
            val processedBitmap =
                Bitmap.createBitmap(mat.width(), mat.height(), Bitmap.Config.ARGB_8888)
            Utils.matToBitmap(mat, processedBitmap)
            return processedBitmap
        }
        return null
    }

}